# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import ujson as json

from coin.base.book.types import TradeSide, create_market_trade_event
from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.korbit_v1.book_builder import BookBuilder


class RestHandlerAsOfWebsocket(WebsocketHandler):
  pass


class KorbitRestHandler(RestHandlerAsOfWebsocket):
  def __init__(self, *, feed_checker_config=None):
    super().__init__(feed_checker_config=feed_checker_config)
    self._last_trade_ids = {}

  def on_book(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict):
      data = json.loads(data.tobytes())
    native_symbol = data['currency_pair']
    if not self.is_subscribed('book', native_symbol):
      return

    book = self._books.setdefault(native_symbol, BookBuilder(native_symbol))
    assert 'bids' in data
    assert 'asks' in data
    bids = data['bids']
    asks = data['asks']
    book.snapshot(bids, asks)
    book.update_timestamp(record.timestamp)
    self.ready = True
    self.publish_event('book', book.symbol, book)

  def on_trade(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict):
      data = json.loads(data.tobytes())
    symbol = data['currency_pair']
    if not self.is_subscribed('trade', symbol):
      return

    self._last_trade_ids.setdefault(symbol, 0)
    last_trade_id = self._last_trade_ids[symbol]

    assert isinstance(data['trades'], list)
    trade_data = reversed(data['trades'])
    tid = None
    for item in trade_data:
      tid = int(item['tid'])
      if tid <= last_trade_id:
        continue
      price = float(item['price'])
      qty = float(item['amount'])
      side = self.infer_side(price, symbol)
      trade_event = create_market_trade_event(record.timestamp, symbol, price, qty, side)
      self.publish_event('trade', symbol, trade_event)

    if tid is not None and tid > last_trade_id:
      self._last_trade_ids[symbol] = tid

  def infer_side(self, price, symbol):
    book = self._books.get(symbol, None)
    if book is not None:
      ask0 = book.ask0().price
      bid0 = book.bid0().price
      if price >= ask0:
        return TradeSide.TRADE_BUY_SIDE
      elif price <= bid0:
        return TradeSide.TRADE_SELL_SIDE
      else:
        return TradeSide.TRADE_NO_SIDE
    else:
      return TradeSide.TRADE_NO_SIDE
